package com.web.crawler.k_means;

import com.web.crawler.compare.SimilarityCalculator;

import java.util.*;
import java.util.stream.Collectors;

public class KMeansClusterer {
    private int k;
    private int maxIterations;
    private List<double[]> dataPoints;
    private List<String> documentNames;
    private List<String> documentContents;
    private Random random;

    public KMeansClusterer(int k, int maxIterations,
                           List<double[]> dataPoints,
                           List<String> documentNames,
                           List<String> documentContents) {
        this.k = k;
        this.maxIterations = maxIterations;
        this.dataPoints = dataPoints;
        this.documentNames = documentNames;
        this.documentContents = documentContents;
        this.random = new Random();
    }

    public List<Cluster> cluster() {
        // 1. 随机初始化聚类中心
        List<double[]> centroids = initializeCentroids();

        List<Cluster> clusters = null;
        boolean converged = false;
        int iteration = 0;
        // 迭代，最多迭代maxIterations次
        while (!converged && iteration < maxIterations) {
            // 2. 分配点到最近的聚类中心
            clusters = assignPointsToClusters(centroids);
            // 3. 重新计算聚类中心
            List<double[]> newCentroids = calculateNewCentroids(clusters);
            // 4. 检查是否收敛
            converged = centroidsConverged(centroids, newCentroids);

            centroids = newCentroids;
            iteration++;
        }

        // 为每个聚类找到代表性文档
        findRepresentativeDocuments(clusters);

        return clusters;
    }

    private List<double[]> initializeCentroids() {
        List<double[]> centroids = new ArrayList<>();
        Set<Integer> chosenIndices = new HashSet<>();

        while (centroids.size() < k) {
            int randomIndex = random.nextInt(dataPoints.size());
            if (!chosenIndices.contains(randomIndex)) {
                centroids.add(Arrays.copyOf(dataPoints.get(randomIndex),
                        dataPoints.get(randomIndex).length));
                chosenIndices.add(randomIndex);
            }
        }

        return centroids;
    }

    private List<Cluster> assignPointsToClusters(List<double[]> centroids) {
        List<Cluster> clusters = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            clusters.add(new Cluster(i, centroids.get(i)));
        }

        for (int i = 0; i < dataPoints.size(); i++) {
            double[] point = dataPoints.get(i);
            int closestClusterId = findClosestCluster(point, centroids);
            clusters.get(closestClusterId).addMember(i, point);
        }

        return clusters;
    }

    private int findClosestCluster(double[] point, List<double[]> centroids) {
        int closestClusterId = 0;
        double minDistance = Double.MAX_VALUE;

        // 循环遍历，找到最近的聚类
        for (int i = 0; i < centroids.size(); i++) {
            // 计算距离
            double distance = SimilarityCalculator.cosineDistance(point, centroids.get(i));
            if (distance < minDistance) {
                minDistance = distance;
                closestClusterId = i;
            }
        }

        return closestClusterId;
    }

    private List<double[]> calculateNewCentroids(List<Cluster> clusters) {
        List<double[]> newCentroids = new ArrayList<>();

        for (Cluster cluster : clusters) {
            if (cluster.getMembers().isEmpty()) {
                // 如果聚类为空，保持原中心
                newCentroids.add(cluster.getCentroid());
                continue;
            }

            double[] newCentroid = new double[dataPoints.get(0).length];

            // 计算成员点的平均值
            for (Integer memberIndex : cluster.getMembers()) {
                double[] memberPoint = dataPoints.get(memberIndex);
                for (int i = 0; i < newCentroid.length; i++) {
                    newCentroid[i] += memberPoint[i];
                }
            }

            for (int i = 0; i < newCentroid.length; i++) {
                newCentroid[i] /= cluster.getMembers().size();
            }

            newCentroids.add(newCentroid);
        }

        return newCentroids;
    }

    private boolean centroidsConverged(List<double[]> oldCentroids,
                                       List<double[]> newCentroids) {
        double threshold = 0.001;

        for (int i = 0; i < oldCentroids.size(); i++) {
            double distance = SimilarityCalculator.cosineDistance(
                    oldCentroids.get(i), newCentroids.get(i));
            if (distance > threshold) {
                return false;
            }
        }

        return true;
    }

    private void findRepresentativeDocuments(List<Cluster> clusters) {
        for (Cluster cluster : clusters) {
            List<Integer> members = cluster.getMembers();
            if (members.isEmpty()) continue;

            // 计算每个成员点到中心的距离
            List<DocumentDistance> documentDistances = new ArrayList<>();
            for (int memberIndex : members) {
                double distance = SimilarityCalculator.cosineDistance(
                        dataPoints.get(memberIndex), cluster.getCentroid());
                documentDistances.add(
                        new DocumentDistance(memberIndex, distance));
            }

            // 按距离排序，取最近的5个作为代表性文档
            documentDistances.sort(Comparator.comparingDouble(DocumentDistance::getDistance));
            List<Integer> representativeIndices = documentDistances.stream()
                    .limit(5)
                    .map(DocumentDistance::getDocumentIndex)
                    .collect(Collectors.toList());

            cluster.setRepresentativeDocuments(representativeIndices);
        }
    }
}

class Cluster {
    private int id;
    private double[] centroid;
    private List<Integer> members;
    private List<Integer> representativeDocuments;

    public Cluster(int id, double[] centroid) {
        this.id = id;
        this.centroid = centroid;
        this.members = new ArrayList<>();
    }

    public void addMember(int documentIndex, double[] point) {
        members.add(documentIndex);
    }

    // Getters and setters
    public int getId() { return id; }
    public double[] getCentroid() { return centroid; }
    public List<Integer> getMembers() { return members; }
    public List<Integer> getRepresentativeDocuments() { return representativeDocuments; }
    public void setRepresentativeDocuments(List<Integer> representativeDocuments) {
        this.representativeDocuments = representativeDocuments;
    }
}

class DocumentDistance {
    private int documentIndex;
    private double distance;

    public DocumentDistance(int documentIndex, double distance) {
        this.documentIndex = documentIndex;
        this.distance = distance;
    }

    public int getDocumentIndex() { return documentIndex; }
    public double getDistance() { return distance; }
}